//
// Header file for hprim.
//
// Copyright (C) 2009-2011 Indrek Mandre <indrek(at)mare.ee>
// http://www.mare.ee/indrek/perc2/, http://code.google.com/p/perc2/
// 
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
// 
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
// 
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
// 

#ifndef __hprim_hpp__
#define __hprim_hpp__

#include <vector>

#include <math.h>
#include <assert.h>

#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

#define DEFAULT_BASEPOW 2
#define DEFAULT_LAYERSTART 0
#define DEFAULT_LAYEREND 50

#define HF_FAST
//#define HF_SMOOTH
//#define HF_FLAT
//#define HF_CONE

//typedef double VALUE_TYPE;
typedef long double VALUE_TYPE;

inline double v2d (VALUE_TYPE v)
{
  return (double)v;
}

class RNG
{
public:
  RNG ()
  {
    //r = gsl_rng_alloc (gsl_rng_ranlxd2);
    r = gsl_rng_alloc (gsl_rng_mt19937);
  }

  ~RNG ()
  {
    gsl_rng_free (r);
  }

  void set_seed (unsigned long int seed)
  {
    gsl_rng_set (r, seed);
  }

  void init (const gsl_rng_type *type, unsigned long int seed)
  {
    printf ("RNG::init: type=%s, seed=%lu\n", type->name, seed);
    gsl_rng_free (r);
    r = gsl_rng_alloc (type);
    set_seed (seed);
  }

  inline double get ()
  {
    return gsl_rng_uniform (r) - 0.5;
  }

  inline double gaussian ()
  {
    return gsl_ran_gaussian (r, 1.0);
  }

  gsl_rng *r;
};

extern RNG rng;

class IncrementLayer
{
public:
  IncrementLayer (VALUE_TYPE cell_width, VALUE_TYPE height, VALUE_TYPE step)
      : cell_width(cell_width), height(height)
  {
    xinc = step / cell_width;
    printf ("Layer: %g %g %g maxampl=%g\n", v2d(cell_width), v2d(height), v2d(step), v2d(xinc * height));
    reset();
  }

  void reset ()
  {
    pv = rng.get() * height;
    mv = rng.get() * height;
    nv = rng.get() * height;
    do {
      x = rng.get() + 0.5;
    } while ( x < 0 || x >= 1.0 );
  }

  VALUE_TYPE next()
  {
    x += xinc;
    while ( x >= 1.0 )
      {
        x -= 1.0;
        pv = mv;
        mv = nv;
        nv = rng.get() * height;
      }

    VALUE_TYPE fp = (1.5 * 1.5 - (x + 0.5) * (x + 0.5)) / (1.5 * 1.5);
    VALUE_TYPE fm = (1.5 * 1.5 - (0.5 - x) * (0.5 - x)) / (1.5 * 1.5);
    VALUE_TYPE fn = (1.5 * 1.5 - (1.5 - x) * (1.5 - x)) / (1.5 * 1.5);

    VALUE_TYPE ret = 0;

#ifdef HF_FAST
    ret += fp * pv;
    ret += fm * mv;
    ret += fn * nv;
#elif defined(HF_SMOOTH)
    ret += fp * fp * fp * fp * pv;
    ret += fm * fm * fm * fm * mv;
    ret += fn * fn * fn * fn * nv;
#else
    #error "TODO: not supported or unknown layer profile"
#endif
    return ret;
  }

private:
  VALUE_TYPE field_width, cell_width, height;
  VALUE_TYPE pv, mv, nv;
  VALUE_TYPE x, xinc;
};

/** Sequential access fBm */
class IncrementField
{
public:
  IncrementField (VALUE_TYPE increment = 1.0, VALUE_TYPE H = 0.5, VALUE_TYPE gbase = 2.0, int layer_start = 0, int layer_end = 40)
      : increment(increment), H(H), gbase(gbase)
  {
    for ( int i = layer_start; i <= layer_end; i++ )
      {
        VALUE_TYPE cell_width = pow (gbase, i);
        VALUE_TYPE height = pow (gbase, H * i);
        layers.push_back (new IncrementLayer (cell_width, height, increment));
      }
  }

  ~IncrementField ()
  {
    for ( layers_t::iterator it = layers.begin(); it != layers.end(); it++ )
        delete *it;
  }

  VALUE_TYPE next ()
  {
    VALUE_TYPE ret = 0;
    for ( size_t i = 0; i < layers.size(); i++ )
        ret += layers[i]->next();
    return ret;
  }

  void reset ()
  {
    for ( size_t i = 0; i < layers.size(); i++ )
        layers[i]->reset();
  }

private:
  typedef std::vector<IncrementLayer *> layers_t;
  layers_t layers;
  VALUE_TYPE increment, H;
  VALUE_TYPE gbase;
};

class Layer
{
public:
  Layer (VALUE_TYPE field_width, VALUE_TYPE cell_width, VALUE_TYPE height)
      : field_width(field_width), cell_width(cell_width), height(height)
  {
    printf ("Layer width=%g, height=%g\n", v2d(cell_width), v2d(height));
    half_cell_width = 0.5 * cell_width;
    cellcount = field_width/cell_width + 5;
    cells = new double [cellcount];
    ohcw = 1.5 * cell_width;
    recp_ohcw = 1.0 / ohcw;
    ohcw2 = ohcw * ohcw;
    recp_ohcw2 = 1.0 / ohcw2;
    reset();
  };

  ~Layer ()
  {
    delete [] cells;
  }

  void reset ()
  {
    for ( size_t i = 0; i < cellcount; i++ )
        cells[i] = rng.get() * height;

    offset = rng.get() * cell_width;

#if 0
    for ( size_t i = 0; i < cellcount; i++ )
        cells[i] = 0;
    cells[4] = 1;
    cells[5] = 1;
    cells[9] = 1;
#endif
  }

  VALUE_TYPE get (VALUE_TYPE x) const
  {
    x += offset;
    int m = (int)(x / cell_width);
#ifdef HF_FAST
    VALUE_TYPE px = (m - 1) * cell_width + half_cell_width;
    VALUE_TYPE mx = m * cell_width + half_cell_width;
    VALUE_TYPE nx = (m + 1) * cell_width + half_cell_width;
    VALUE_TYPE fp = recp_ohcw2 * (ohcw2 - (px - x) * (px - x));
    VALUE_TYPE fm = recp_ohcw2 * (ohcw2 - (mx - x) * (mx - x));
    VALUE_TYPE fn = recp_ohcw2 * (ohcw2 - (nx - x) * (nx - x));
    VALUE_TYPE ret = 0;
    if ( fp > 0 ) ret += fp * cells[2 + m - 1];
    if ( fm > 0 ) ret += fm * cells[2 + m];
    if ( fn > 0 ) ret += fn * cells[2 + m + 1];
    return ret;
#elif defined(HF_CONE)
    VALUE_TYPE px = (m - 1) * cell_width + half_cell_width;
    VALUE_TYPE mx = m * cell_width + half_cell_width;
    VALUE_TYPE nx = (m + 1) * cell_width + half_cell_width;

    VALUE_TYPE fp = recp_ohcw * (ohcw - fabs(px - x));
    VALUE_TYPE fm = recp_ohcw * (ohcw - fabs(mx - x));
    VALUE_TYPE fn = recp_ohcw * (ohcw - fabs(nx - x));
    VALUE_TYPE ret = 0;
    if ( fp > 0 ) ret += fp * cells[2 + m - 1];
    if ( fm > 0 ) ret += fm * cells[2 + m];
    if ( fn > 0 ) ret += fn * cells[2 + m + 1];
    return ret;
#elif defined(HF_SMOOTH)
    VALUE_TYPE px = (m - 1) * cell_width + half_cell_width;
    VALUE_TYPE mx = m * cell_width + half_cell_width;
    VALUE_TYPE nx = (m + 1) * cell_width + half_cell_width;

    VALUE_TYPE fp = recp_ohcw2 * (ohcw2 - (px - x) * (px - x));
    VALUE_TYPE fm = recp_ohcw2 * (ohcw2 - (mx - x) * (mx - x));
    VALUE_TYPE fn = recp_ohcw2 * (ohcw2 - (nx - x) * (nx - x));
    VALUE_TYPE ret = 0;
    if ( fp > 0 ) ret += fp * fp * fp * fp * cells[2 + m - 1];
    if ( fm > 0 ) ret += fm * fm * fm * fm * cells[2 + m];
    if ( fn > 0 ) ret += fn * fn * fn * fn * cells[2 + m + 1];
    return ret;
#elif defined(HF_FLAT)
    return cells[2 + m];
#else
    #error "unkknown layer profile"
#endif
  }

private:
  VALUE_TYPE offset;
  VALUE_TYPE field_width, cell_width, height;
  double *cells;
  size_t cellcount;
  VALUE_TYPE half_cell_width;
  VALUE_TYPE ohcw, ohcw2, recp_ohcw, recp_ohcw2;
};

/** Random access fBm */
class Field
{
public:
  Field (VALUE_TYPE field_width, VALUE_TYPE H = 0.5, VALUE_TYPE gbase = 2.0, size_t layer_start = 0, size_t layer_end = 40)
      : field_width (field_width), H(H), gbase(gbase)
  {
    for ( size_t i = layer_start; i <= layer_end; i++ )
      {
        VALUE_TYPE cell_width = pow (gbase, i);
        VALUE_TYPE height = pow (gbase, H * i);
        layers.push_back (new Layer (field_width, cell_width, height));
      }
  }

  ~Field ()
  {
    for ( layers_t::iterator it = layers.begin(); it != layers.end(); it++ )
        delete *it;
  }

  VALUE_TYPE get (VALUE_TYPE x) const
  {
    VALUE_TYPE ret = 0;
    for ( size_t i = 0; i < layers.size(); i++ )
        ret += layers[i]->get(x);
    return ret;
  }

  inline double get_width () const { return field_width; }

  void reset ()
  {
    for ( size_t i = 0; i < layers.size(); i++ )
        layers[i]->reset();
  }

private:
  typedef std::vector<Layer *> layers_t;
  layers_t layers;
  VALUE_TYPE field_width;
  VALUE_TYPE H;
  VALUE_TYPE gbase;
};

#define MAX_DATA_POINTS 20
/** Base class for sequential access based H estimation. */
class IHSAMPLER
{
public:
  IHSAMPLER (size_t n) : n(n)
  {
    samples = new VALUE_TYPE [n + 1];
    data = new long double [MAX_DATA_POINTS];
    data2 = new long double [MAX_DATA_POINTS];
    data_count = new size_t [MAX_DATA_POINTS];
    for ( size_t i = 0; i < MAX_DATA_POINTS; i++ )
      {
        data[i] = 0;
        data2[i] = 0;
        data_count[i] = 0;
      }
  }

  void print (const char *path)
  {
    FILE *fp = fopen (path, "w");
    assert (fp);
    for ( size_t i = 0; i < MAX_DATA_POINTS; i++ )
      {
        if ( data_count[i] )
            fprintf (fp, "%d,1,%zu,%.33Lg,%.33Lg\n", 1 << i, data_count[i], data[i], data2[i]);
      }
    fclose (fp);
  }

  virtual ~IHSAMPLER ()
  {
    delete [] samples;
    delete [] data;
    delete [] data2;
    delete [] data_count;
  }

  virtual VALUE_TYPE next () = 0;

  void sample ()
  {
    for ( size_t i = 0; i <= n; i++ )
        samples[i] = next();

    for ( size_t i = 0; i < MAX_DATA_POINTS; i++ )
      {
        for ( size_t x = 0; x + (1 << i) <= n; x++ )
          {
            VALUE_TYPE v1 = samples[x];
            VALUE_TYPE v2 = samples[x+(1<<i)];
            if ( v1 == 0.0 || v2 == 0.0 )
                continue;
            VALUE_TYPE diff = v1 - v2;
            VALUE_TYPE diff2 = diff * diff;
            VALUE_TYPE diff4 = diff2 * diff2;
            data[i] += diff2;
            data2[i] += diff4;
            data_count[i]++;
          }
      }
  }

private:
  size_t n;
  VALUE_TYPE *samples;
  size_t datapoints;
  long double *data;
  long double *data2;
  size_t *data_count;
};

/** Base class for random access based H estimation. */
class HSAMPLER
{
public:
  HSAMPLER(VALUE_TYPE step) : step(step)
  {
    data = new long double [MAX_DATA_POINTS];
    data2 = new long double [MAX_DATA_POINTS];
    data_count = new size_t [MAX_DATA_POINTS];
    for ( size_t i = 0; i < MAX_DATA_POINTS; i++ )
      {
        data[i] = 0;
        data2[i] = 0;
        data_count[i] = 0;
      }
  }

  void print (const char *path)
  {
    FILE *fp = fopen (path, "w");
    assert (fp);
    for ( size_t i = 0; i < MAX_DATA_POINTS; i++ )
        fprintf (fp, "%g,1,%zu,%.33Lg,%.33Lg\n", v2d(step * (1 << i)), data_count[i], data[i], data2[i]);
    fclose (fp);
  }

  virtual ~HSAMPLER ()
  {
    delete [] data;
    delete [] data2;
    delete [] data_count;
  }

  virtual VALUE_TYPE get (VALUE_TYPE x) = 0;

  void sample (VALUE_TYPE a, VALUE_TYPE b)
  {
    n = (b - a) / step;
    VALUE_TYPE *samples = new VALUE_TYPE [n + 1];
    printf ("Alloced samples %zu\n", n + 1);
    for ( size_t i = 0; i <= n; i++ )
      {
        VALUE_TYPE v = get (a + step * i);
        samples[i] = v;
      }

    size_t begin, end;
    begin = 0;
    end = n;

    printf ("Working in the range [%zu,%zu]/%zu\n", begin, end, end - begin);

    for ( size_t i = 0; i < MAX_DATA_POINTS; i++ )
      {
        for ( size_t x = begin; x + (1 << i) <= end; x++ )
          {
            VALUE_TYPE v1 = samples[x];
            VALUE_TYPE v2 = samples[x+(1<<i)];
            if ( v1 == 0.0 || v2 == 0.0 )
                continue;
            VALUE_TYPE diff = v1 - v2;
            VALUE_TYPE diff2 = diff * diff;
            VALUE_TYPE diff4 = diff2 * diff2;
            data[i] += diff2;
            data2[i] += diff4;
            data_count[i]++;
            //printf ("SAmple! %g %g, %g; data[%zu]=%g\n", v2d(v1), v2d(v2), v2d(diff), i, v2d(data[i]));
          }
      }

    //printf ("Done...\n");
    //getchar();

    delete [] samples;
  }

private:
  VALUE_TYPE step;
  size_t n;
  long double *data;
  long double *data2;
  size_t *data_count;
};

/** Sampler used by testh for H estimation. */
class testsampler : public IHSAMPLER
{
public:
  testsampler (IncrementField& f, size_t n) : IHSAMPLER (n), f(f)
  {
  }

  VALUE_TYPE next ()
  {
    return f.next ();
  }

private:
  IncrementField& f;
};

/** FBm zero-set size counter. */
class zcounter
{
  // XXX: cache size should depend on H, as it gets smaller
  // so should cache, otherwise we get excessive overlap and
  // things will instead get slower
  #define ZCACHE_SIZE 16384
  //#define ZCACHE_SIZE 16
  #define ZCACHE_MID (ZCACHE_SIZE/2)
public:
  zcounter (IncrementField& f, size_t n) : f(f), n(n)
  {
    printf ("zcounter: %zu slots\n", n + 1);
    samplesd = new double [n + 1];
  }

  ~zcounter ()
  {
    delete [] samplesd;
  }

  void sample ()
  {
    release_cache();

    avg = f.next();
    mx = mn = 0;
    printf ("Sampling [%zu, %zu]\n", (size_t)0, n);
    size_t di = 0;
    for ( size_t i = 0; i <= n; i++, di++ )
      {
        double v = f.next() - avg;
        mx = std::max (mx, v);
        mn = std::min (mn, v);
        if ( i > 1 )
          {
            if ( samplesd[di - 1] > samplesd[di - 2] && v >= samplesd[di - 1] )
                di--;
            else if ( samplesd[di - 1] < samplesd[di - 2] && v <= samplesd[di - 1] )
                di--;
          }
        samplesd[di] = v;
      }
    samplesd_range = di;

    avg = 0.5 * (mx + mn);

    build_cache();
  }

  size_t get (double z)
  {
    int p = (z - avg) / cache_step + ZCACHE_MID;
    if ( p < 0 )
        p = 0;
    if ( p >= ZCACHE_SIZE )
        p = ZCACHE_SIZE - 1;

    std::vector<int>& c = cache[p];
    size_t ret = 0;
    for ( size_t i = 0; i < c.size(); i++ )
        ret += (z < samplesd[c[i]]) ^ (z < samplesd[c[i] + 1]);
    return ret;
  }

  inline std::pair<double, double> get_range () const { return std::make_pair (mn, mx); }
  
  void build_cache ()
  {
    double range = mx - mn;
    printf ("build_cache(), range=%g, avg=%g..\n", range, avg);
    for ( size_t i = 0; i < ZCACHE_SIZE; i++ )
        cache[i].clear();

    cache_step = range / ZCACHE_SIZE;

    for ( size_t i = 0; i < samplesd_range - 1; i++ )
      {
        int p1 = (samplesd[i] - avg) / cache_step + ZCACHE_MID;
        int p2 = (samplesd[i + 1] - avg) / cache_step + ZCACHE_MID;

        if ( p1 < 0 )
            p1 = 0;
        else if ( p1 >= ZCACHE_SIZE )
            p1 = ZCACHE_SIZE - 1;
        if ( p2 < 0 )
            p2 = 0;
        else if ( p2 >= ZCACHE_SIZE )
            p2 = ZCACHE_SIZE - 1;
        
        if ( p2 < p1 )
            std::swap (p1, p2);
        for ( int j = p1; j <= p2; j++ )
            cache[j].push_back (i);
      }

    printf ("Compression: %zu/%zu of space used up\n", samplesd_range, n + 1);

    //for ( size_t i = 0; i < ZCACHE_SIZE; i++ )
    //    printf ("%zu: %zu\n", i, cache[i].size());
  }

  void release_cache ()
  {
    for ( size_t i = 0; i < ZCACHE_SIZE; i++ )
      {
        std::vector<int> cl;
        cl.swap (cache[i]);
      }
  }

private:
  IncrementField& f;
  size_t n;
  double *samplesd;
  size_t samplesd_range;
  std::vector<int> cache[ZCACHE_SIZE];
  double cache_step;
  double avg;
  double mn, mx;
};

/** Sampler used by hprim for H' estimation. */
class zsampler : public HSAMPLER
{
public:
  //zsampler (zcounter& zc, VALUE_TYPE interval) : HSAMPLER (zc.get_range().first, zc.get_range().second, interval), zc(zc)
  zsampler (zcounter& zc, VALUE_TYPE interval) : HSAMPLER (interval), zc(zc)
  {
  }

  VALUE_TYPE get (VALUE_TYPE x)
  {
    return zc.get (x);
  }

private:
  zcounter& zc;
};

#endif // !__hprim_hpp__

